import torch
import torch.nn as nn
import math

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        """
        类的初始化函数
        d_model: 指词嵌入的维度
        vocab: 指词表的大小
        """
        super(Embeddings, self).__init__()
        # 之后就是调用nn中的预定义层Embedding，获得一个词嵌入对象self.lut
        self.lut = nn.Embedding(vocab, d_model)
        # 最后就是将d_model传入类中
        self.d_model = d_model

    def forward(self, x):
        """
        Embedding层的前向传播器
        参数x: 这里代表输入给模型的单词文本通过词表映射后的index向量
        将文化给self.lut并且根据self.d_model相乘为结果返回
        """
        embedds = self.lut(x)
        return embedds * math.sqrt(self.d_model)

# 测试代码
if __name__ == "__main__":
    # 检查CUDA是否可用
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    voc_dict = {"I": 1, "want": 2, "a": 3, "plane": 4, "to": 5, "travel": 6, "Beijing": 7}
    d_model = 512
    vocab = 10

    # 创建模型并移动到GPU
    embeddings = Embeddings(d_model, vocab).to(device)

    # 创建输入并移动到GPU
    inputs = torch.LongTensor([[1, 2, 3, 4, 0, 6], [1, 2, 5, 6, 5, 7]]).to(device)

    # 前向传播
    outputs = embeddings(inputs)

    print(f"Input shape: {inputs.shape}")
    print(f"Output shape: {outputs.shape}")
    print(f"Input device: {inputs.device}")
    print(f"Output device: {outputs.device}")

